/*
 * @Author: LVGRAPE
 * @LastEditors: LVGRAPE
 */
#include "kalman_2d.h"
#include "rtthread.h"

#define MATRIX_2X2_ADD(A, B, C)  \
    C[0][0] = A[0][0] + B[0][0]; \
    C[0][1] = A[0][1] + B[0][1]; \
    C[1][0] = A[1][0] + B[1][0]; \
    C[1][1] = A[1][1] + B[1][1];
#define MATRIX_2X2_SUB(A, B, C)  \
    C[0][0] = A[0][0] - B[0][0]; \
    C[0][1] = A[0][1] - B[0][1]; \
    C[1][0] = A[1][0] - B[1][0]; \
    C[1][1] = A[1][1] - B[1][1];
#define MATRIX_2X2_MUL(A, B, C)                      \
    C[0][0] = A[0][0] * B[0][0] + A[0][1] * B[1][0]; \
    C[0][1] = A[0][0] * B[0][1] + A[0][1] * B[1][1]; \
    C[1][0] = A[1][0] * B[0][0] + A[1][1] * B[1][0]; \
    C[1][1] = A[1][0] * B[0][1] + A[1][1] * B[1][1];
#define MATRIX_2X2_TRAN(A, B) \
    B[0][0] = A[0][0];        \
    B[0][1] = A[1][0];        \
    B[1][0] = A[0][1];        \
    B[1][1] = A[1][1];


int matrix_2x2_inv(float A[][2], float B[][2])
{
    float div = A[0][0] * A[1][1] - A[0][1] * A[1][0];
    if (div == 0)
    {
        return 0;
    }
    B[0][0] = A[1][1] / div;
    B[0][1] = -A[0][1] / div;
    B[1][0] = -A[1][0] / div;
    B[1][1] = A[0][0] / div;
    return 0;
}

/**
 * @brief 卡尔曼滤波器初始化为速度和位置方式
 * noise up, trust down
 * eg: if(q < r) means the predict is more accurate than measure
 * @param k
 * @param q1 position predict noise
 * @param q2 velocity predict noise
 * @param r1 position measure noise
 * @param r2 velocity measure noise
 * @param dt
 */
void kalman_2d_init_vel_pos(kalman_2d_t *k, float q1, float q2, float r1, float r2, float dt)
{
    memset(k, 0, sizeof(kalman_2d_t));
    /**
     * @brief about transform matrix.Means how last state transfer to current state
     * ag: x1 for true position, x2 for true velocity
     *     z1 for measure position, z2 for measure velocity
     *     w1 for position noise, w2 for velocity noise
     *     v1 for postion measure noise, v2 for velocity measure noise
     *  therefor  x1 = [1] * x1_last + [dt] * x2_last + w1   ---> A = |1 dt| if (dt = 1) A= |1 1|
     *            x2 = [0] * x1_last + [1] * x2_last + w2    --->     |0 1 |                |0 1|
     *     z1 = [1] * x1 + [0] * x2 + v1  ----------------------> H = |1 0|
     *     z2 = [0] * x1 + [1] * x2 + v2  ---------------------->     |0 1|
     *
     */
    k->Q[0][0] = q1;
    k->Q[1][1] = q2;
    k->R[0][0] = r1;
    k->R[1][1] = r2;

    k->A[0][0] = 1;
    k->A[0][1] = 1;//dt=1
    k->A[1][1] = 1;

    k->H[0][0] = 1;
    k->H[1][1] = 1;

    k->I[0][0] = 1;
    k->I[1][1] = 1;

    k->x_hat[0][1] = 1;

    k->P_predict[0][0] = 1;
    k->P_predict[1][1] = 1;
}

/**
 * @brief
 *
 * @param k
 * @param z1 z1 测量位置
 * @param z2 z2 测量速度
 */
void kalman_2d_update_vel_pos(kalman_2d_t *k, float position, float velocity)
{
    float temp1[2][2];
    float temp2[2][2];
    float temp3[2][2];

    k->z[0][0] = position;
    k->z[0][1] = velocity;

    // x_hat_predict = A*x_hat_last
    MATRIX_2X2_MUL(k->A, k->x_hat_last, k->x_hat_predict);
    // rt_kprintf("x_hat_predict: %d %d\t", (int)(k->x_hat_predict[0][0]*1000), (int)(k->x_hat_predict[0][1]*1000));
    // P_predict = A*P_predict_last*A' + Q
    MATRIX_2X2_MUL(k->A, k->P_predict_last, temp1); // A*P_predict_last
    MATRIX_2X2_TRAN(k->A, temp2);                   // A'
    MATRIX_2X2_MUL(temp1, temp2, temp3);            // A*P_predict_last*A'
    MATRIX_2X2_ADD(temp3, k->Q, k->P_predict);      // A*P_predict_last*A'+Q
    // rt_kprintf("P_predict: %d %d\t", (int)(k->P_predict[0][0]*1000), (int)(k->P_predict[0][1]*1000));
    // K = P_predict*H'/(H*P_predict_last*H'+R)
    MATRIX_2X2_TRAN(k->H, temp1);                   // H'
    MATRIX_2X2_MUL(k->H, k->P_predict_last, temp2); // H*P_predict_last
    MATRIX_2X2_MUL(temp1, temp2, temp3);            // H*P_predict_last*H'
    MATRIX_2X2_ADD(temp3, k->R, temp2);             // H*P_predict_last*H'+R
    matrix_2x2_inv(temp2, temp3);                   // 1/(H*P_predict_last*H'+R)
    MATRIX_2X2_MUL(k->P_predict, temp1, temp2);     // P_predict*H'
    MATRIX_2X2_MUL(temp2, temp3, k->K);             // P_predict*H'/(H*P_predict_last*H'+R)
    // rt_kprintf("K: %d %d\t", (int)(k->K[0][0]*1000), (int)(k->K[0][1]*1000));
    // x_hat = x_hat_predict + K*(z-H*x_hat_predict)
    MATRIX_2X2_MUL(k->H, k->x_hat_predict, temp1);     // H*x_hat_predict
    MATRIX_2X2_SUB(k->z, temp1, temp2);                //(z-H*x_hat_predict)
    MATRIX_2X2_MUL(k->K, temp2, temp1);                // K*(z-H*x_hat_predict)
    MATRIX_2X2_ADD(k->x_hat_predict, temp1, k->x_hat); // x_hat_predict + K*(z-H*x_hat_predict)
    // rt_kprintf("x_hat: %d %d\t", (int)(k->x_hat[0][0]*1000), (int)(k->x_hat[0][1]*1000));
    // P = (I-K*H)*P_predict
    MATRIX_2X2_MUL(k->K, k->H, temp1);         // K*H
    MATRIX_2X2_SUB(k->I, temp1, temp2);        //(I-K*H)
    MATRIX_2X2_MUL(temp2, k->P_predict, k->P); //(I-K*H)*P_predict
    // rt_kprintf("P: %d %d\n", (int)(k->P[0][0]*1000), (int)(k->P[0][1]*1000));

    k->x_hat_last[0][0] = k->x_hat[0][0];
    k->x_hat_last[0][1] = k->x_hat[0][1];
    k->x_hat_last[1][0] = k->x_hat[1][0];
    k->x_hat_last[1][1] = k->x_hat[1][1];

    k->P_predict_last[0][0] = k->P_predict[0][0];
    k->P_predict_last[0][1] = k->P_predict[0][1];
    k->P_predict_last[1][0] = k->P_predict[1][0];
    k->P_predict_last[1][1] = k->P_predict[1][1];
}

#if 0
#include <stdio.h>
#include <stdlib.h>
int main()
{
    kalman_2d_t k;
    kalman_2d_init_vel_pos(&k, 0.1, 0.1);
    float x1, x2, z1, z2;
    printf("实际位置,实际速度,测量位置,测量速度,先验位置,先验速度,后验位置,后验速度\n");
    for (uint16_t i = 0; i < 200; i++)
    {
        x2 = (rand() % 10) / 20.f + (rand() % 10) / 50.f;//速度
        x1 = x2 + x1 + (rand() % 10) / 50.f;//位置
        z1 = x1 + (rand() % 10) / 50.f;//测量位置
        z2 = x2 + (rand() % 10) / 50.f;//测量速度
        kalman_2d_update(&k,z1, z2);
        printf("%.2f, %.2f, %.2f, %.2f, %.2f, %.2f, %.2f, %.2f\n", x1, x2, z1, z2, k.x_hat_predict[0], k.x_hat_predict[1], k.x_hat[0], k.x_hat[1]);

    }
    return 0;
}
#endif
